import networkx as nx
# import matplotlib.pyplot as plt
import os
import pandas as pd
import Link_Prediction_Scores_All as lp
import pickle
import os
import numpy as np
import embMethod_simulation.Node2Vec as node2vec
import Link_Prediction_Scores_All
import figure
import matplotlib.pyplot as plt

# import tensorflow as tf


NUM_REPEATS = 10
F_Num_REPEATS = NUM_REPEATS * 1.0
RANDOM_SEED = 0
FRAC_EDGES_HIDDEN = [0.1]
dimen = [0.1]

DEAL_NAME = ['grid_2d-100-100']
method = 'Line'

for deal_name in DEAL_NAME:
    res_sum = pd.DataFrame()
    combined_dir = 'D:/data-processed/{}-adj.pkl'.format(deal_name)
    with open(combined_dir, 'rb') as f:
        data = pickle.load(f)
    if('grid' in deal_name):
      adj,news_edges = data
    else:
      adj = data
    test_frac = 0.1
    TRAIN_TEST_SPLITS_FOLDER = 'D:/data/splited/ACT13_sample/1_1/{}/'.format(deal_name)

    AUCs = []
    pres = []
    weak_pres = []

    p = 0.5
    q = 2

    tmp_AUCs = []
    tmp_pres = []
    tmp_weak_pres = []

    for j in range(NUM_REPEATS):
        print('第{}次执行-原始图嵌入方法'.format(j))
        experiment_name = '{}-{}-{}-hidden'.format(deal_name, test_frac, j)
        train_test_split_file = TRAIN_TEST_SPLITS_FOLDER + experiment_name + '.pkl'

        AUC, pre, weak_pre = lp.calculate_scores(adj, test_frac=test_frac, val_frac=0,
                                                 random_state=RANDOM_SEED, verbose=2,
                                                 train_test_split_file=train_test_split_file, multip=0.6, dims=10, k=2,
                                                 diff=1, method=method,node2vec_dim=128)

        tmp_AUCs.append(AUC)
        tmp_pres.append(pre)
        tmp_weak_pres.append(weak_pre)

    tmp_AUC = np.array(tmp_AUCs).sum(axis=0) / F_Num_REPEATS
    tmp_pre = np.array(tmp_pres).sum(axis=0) / F_Num_REPEATS
    tmp_weak_pre = np.array(tmp_weak_pres).sum(axis=0) / F_Num_REPEATS

    AUC = [float('{:.4f}'.format(i)) for i in tmp_AUC]
    pre = [float('{:.4f}'.format(i)) for i in tmp_pre]
    weak_pre = [float('{:.4f}'.format(i)) for i in tmp_weak_pre]

    AUCs.append(AUC)
    pres.append(pre)
    weak_pres.append(weak_pre)

    name = ['MLP-ori']
    evaluation_name = [AUCs, pres, weak_pres]

    test1 = pd.DataFrame(index=['AUC'],columns=name, data=AUCs)

    test2 = pd.DataFrame(index=['pre'],columns=name, data=pres)
    # test.to_csv('{}-ori-All-pre.csv'.format(deal_name))

    test3 = pd.DataFrame(index=['weak_pre'],columns=name, data=weak_pres)
    # test.to_csv('{}-ori-All-weak_pre.csv'.format(deal_name))

    res_ori = pd.concat([test1,test2,test3],axis=0)
    res_ori.to_csv('{}-{}-ori.csv'.format(method,deal_name))

    AUCs = []
    pres = []
    weak_pres = []
    components = [128,132,135,140,145,150,155,160]
    for component in components:


        tmp_AUCs = []
        tmp_pres = []
        tmp_weak_pres = []
        file_name = 'D:/data/similarity_matrix/'
        for j in range(NUM_REPEATS):
            print('第{}次执行-自定义方法'.format(j))
            experiment_name = '{}-{}-{}-hidden'.format(deal_name, test_frac, j)
            train_test_split_file = TRAIN_TEST_SPLITS_FOLDER + experiment_name + '.pkl'

            train_test_split = None
            with open(train_test_split_file, 'rb') as f:
                train_test_split = pickle.load(f)
                print('Found existing train-test split!')

            adj_train, train_edges, train_edges_false, val_edges, val_edges_false, \
            test_edges, test_edges_false, num_dif_community, train_num_dif_community = train_test_split
            multip = 0.1

            def get_scores(embedding):
                MLP_Scores= Link_Prediction_Scores_All.get_all_scores(
                    adj_train, train_test_split, edge_score_mode="edge-emb",
                    verbose=1, multip=multip, k=2, dims=10, emb_matrix=embedding, Embedding_Method=method)
                AUC = []
                pre = []
                s_pre = []
                AUC.append(MLP_Scores['test_AUC'])
                pre.append(MLP_Scores['test_precision'])
                s_pre.append(MLP_Scores['test_weak_precision'])

                return AUC, pre, s_pre
            if(method=='Node2Vec'):
              oriS_name = file_name+method+'/'+deal_name+'-ori-{}-{}.pkl'.format(p,q)

            # with open(oriS_name, 'rb') as f:
            #   S = pickle.load(f)

            newS_name= file_name+method+'/'+deal_name+'-pca-{}.pkl'.format(component)
            with open(newS_name, 'rb') as f:
                newS = pickle.load(f)
            AUC, pre, weak_pre = get_scores(newS)

            tmp_AUCs.append(AUC)
            tmp_pres.append(pre)
            tmp_weak_pres.append(weak_pre)

        tmp_AUC = np.array(tmp_AUCs).sum(axis=0) / F_Num_REPEATS
        tmp_pre = np.array(tmp_pres).sum(axis=0) / F_Num_REPEATS
        tmp_weak_pre = np.array(tmp_weak_pres).sum(axis=0) / F_Num_REPEATS

        AUC = [float('{:.4f}'.format(i)) for i in tmp_AUC]
        pre = [float('{:.4f}'.format(i)) for i in tmp_pre]
        weak_pre = [float('{:.4f}'.format(i)) for i in tmp_weak_pre]

        AUCs.append(AUC)
        pres.append(pre)
        weak_pres.append(weak_pre)

        evaluation_name = [AUCs,pres,weak_pres]

        # test = pd.DataFrame(index=methods,columns=name, data=evaluation_name)
        # test.to_csv('{}-similarity-All.csv'.format(deal_name))

        # test = pd.DataFrame(index=methods,columns=name, data=pres)
        # test.to_csv('{}-similarity-All-pre.csv'.format(deal_name))
        #
        # test = pd.DataFrame(index=methods,columns=name, data=weak_pres)
        # test.to_csv('{}-similarity-All-weak_pre.csv'.format(deal_name))
    test1 = pd.DataFrame(index=['AUC'], columns=components, data=np.reshape(AUCs,(1,len(components))))

    test2 = pd.DataFrame(index=['pre'], columns=components, data=np.reshape(pres,(1,len(components))))
        # test.to_csv('{}-ori-All-pre.csv'.format(deal_name))

    test3 = pd.DataFrame(index=['weak_pre'], columns=components, data=np.reshape(weak_pres,(1,len(components))))
        # test.to_csv('{}-ori-All-weak_pre.csv'.format(deal_name))
    res_sum = pd.concat([test1, test2, test3], axis=0)

    res_sum.to_csv('{}-{}-pca.csv'.format(method,deal_name))

components = [128,132,135,140,145,150,155,160]
x = components
colors = ['g', 'b', 'y', 'c', 'm','a','b','d','r']
marker = ["o", "x", "^", ".", "+","1","2","3","4","5"]
plt.figure(1)
for deal_name,i in zip(DEAL_NAME,range(len(DEAL_NAME))):
    df_ori = pd.read_csv('{}-{}-ori'.format(method, deal_name) + '.csv')
    df_pca = pd.read_csv('{}-{}-pca'.format(method, deal_name) + '.csv')
    plt.plot(x, df_pca.values[0][1:], label=deal_name, color=colors[i], marker=marker[i])
    plt.plot(x, np.repeat(df_ori.values[0][1:], len(components)), label='ori-' + str(deal_name), color=colors[i])
    # ori_ori_weak_pres = np.repeat(ori_weak_pre, len(weak_pre))
    # plt.plot(x, ori_AUCs, color = 'g',label='ori_AUC')
    # plt.plot(x, ori_pres, color = 'b',label='ori_pre')
    # plt.plot(x, ori_ori_weak_pres, color = 'r',label='ori_weak_pre')
    # plt.plot(x, AUC,label='AUC',color = 'g',marker = "o")
    # plt.plot(x, pre,label='pre',color = 'b',marker = "x")
    # plt.plot(x, weak_pre,label='weak_pre',color = 'r',marker = "^")
plt.title(method  +'-AUC')
plt.xlabel('embedding_size')
plt.legend()
plt.savefig("{}.png".format(method + '-' + 'MLP'+'-AUC'))
plt.show()

plt.figure(2)
for deal_name,i in zip(DEAL_NAME,range(len(DEAL_NAME))):
    df_ori = pd.read_csv('{}-{}-ori'.format(method, deal_name) + '.csv')
    df_pca = pd.read_csv('{}-{}-pca'.format(method, deal_name) + '.csv')
    plt.plot(x, df_pca.values[1][1:], label=deal_name, color=colors[i], marker=marker[i])
    plt.plot(x, df_pca.values[2][1:], label=deal_name+'-weak', color=colors[i], marker=marker[i],linestyle='--')
    plt.plot(x, np.repeat(df_ori.values[1][1:], len(components)), label='ori-' + str(deal_name), color=colors[i])
    plt.plot(x, np.repeat(df_ori.values[2][1:], len(components)), label='ori-' + str(deal_name)+'-weak', color=colors[i],linestyle=':')
    # ori_ori_weak_pres = np.repeat(ori_weak_pre, len(weak_pre))
    # plt.plot(x, ori_AUCs, color = 'g',label='ori_AUC')
    # plt.plot(x, ori_pres, color = 'b',label='ori_pre')
    # plt.plot(x, ori_ori_weak_pres, color = 'r',label='ori_weak_pre')
    # plt.plot(x, AUC,label='AUC',color = 'g',marker = "o")
    # plt.plot(x, pre,label='pre',color = 'b',marker = "x")
    # plt.plot(x, weak_pre,label='weak_pre',color = 'r',marker = "^")
plt.title(method +'-pre+weak_pre')
plt.xlabel('embedding_size')
plt.legend()
plt.savefig("{}.png".format(method + '-' + 'MLP'+'-pre+weak_pre'))
plt.show()



